﻿#include "PyramidTemplateMatch.h"
#include <opencv2/highgui.hpp>
#include <iostream>
#include "iostream_opencv.h"

using namespace std;
namespace Qly {

static const double ANGLE_EPS = 0.001;
PyramidTemplateMatch::PyramidTemplateMatch(int maxDepth)
    :m_maxDepth(maxDepth)
{

}

bool PyramidTemplateMatch::setModelByAngleCount(const cv::Mat &model, double angleMin, double angleMax, size_t k)
{
//    if(k == 0 || k > 12) return false;
    cv::Mat mask = cv::Mat(model.rows, model.cols, CV_8U, 255);
    return setModelByAngleCount(model, mask, angleMin, angleMax, k);
}

// bool PyramidTemplateMatch::setModelByDetAngle(const cv::Mat &model, double angleMin, double angleMax, double detAngle)
// {
//     cv::Mat mask = cv::Mat(model.rows, model.cols, CV_8U, 255);
//     return setModelByDetAngle(model, mask, angleMin, angleMax, detAngle);
// }

bool PyramidTemplateMatch::setModelByAngleCount(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, size_t k)
{
    if(k == 0 || k > 12) return false;
    // size_t angleCount = (1 << k);
    if(angleMax < angleMin) return false;
    m_angleMin = angleMin;
    m_angleMax = angleMax;
    m_angleCount = (1 << k);
    std::cout << "setModelByAngleCount:m_angleCount = " << m_angleCount << endl;
    m_angleDet = (angleMax - angleMin) / m_angleCount;
    return buildModels(model, mask);
}

// bool PyramidTemplateMatch::setModelByDetAngle(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, double detAngle)
// {
//     if(detAngle <= 0) return false;
//     if(angleMax < angleMin) return false;

//     m_angleMin = angleMin;
//     m_angleDet = detAngle;
//     m_angleCount = ( angleMax - angleMin ) / detAngle + 0.9; //
//     m_angleMax = angleMin + detAngle * m_angleCount;
//     std::cout << "angleMin = " << angleMin << ", angleMax = " << angleMax << ", angleCount = " << m_angleCount << ", angleDet =" << detAngle << std::endl;
//     return buildModels(model, mask);
// }


bool PyramidTemplateMatch::match(const cv::Mat &image, cv::RotatedRect &rrt, double &score)
{
    if(image.channels() != 1)
    {
        std::cout << "PyramidTemplateMatch::match() image must have only one channel." << std::endl;
        return false;
    }
    if(m_pyMatch.size() == 0)
    {
        std::cout << "m_pyMatch.size() == 0, return false" << std::endl;
        return false;
    }

    if(m_depthCount == 0) return m_pyMatch[0].rotatedMatch(image, rrt, score);

    cv::Mat tempImage = image.clone();
    std::vector<cv::Mat> imagePyramid;
    imagePyramid.push_back(tempImage.clone());// 第 0 层就是原始图像。
    for(int i = 1; i <= m_depthCount; i++)  // 构造图像金字塔
    {
        cv::pyrDown(tempImage, tempImage);
        imagePyramid.push_back(tempImage.clone());
    }

    size_t depthCount = m_depthCount;
    cv::Rect rect;
    // 第一次是全局搜索。后面都是小区域内搜索
    bool ret = m_pyMatch[depthCount].rotatedMatch(imagePyramid[depthCount], rrt, score);
//    std::cout << "first search, depthCount = " << depthCount << "  rrt = " << rrt << std::endl;

    if(!ret) return false;
    do
    {
        depthCount --;
        rrt.center *= 2;
        m_pyMatch[depthCount].roi(imagePyramid[depthCount], rrt, 3, rect);
        std::cout << "roi = " << rect << std::endl;
        // 到这里就获得了合适的 ROI 区域
        m_pyMatch[depthCount].regionMatch(imagePyramid[depthCount], rect, rrt.angle, rrt, score);
//        std::cout << "depth = " << depthCount <<  ", rrt = " << rrt << ", score = " << score << "   regionMatch" << std::endl;
    }while(depthCount != 0);

    return true;
}

bool PyramidTemplateMatch::match0(const cv::Mat &image, cv::RotatedRect &rrt, double &score)
{
    if(image.channels() != 1)
    {
        std::cerr << "PyramidTemplateMatch::match() image must have only one channel." << std::endl;
        return false;
    }
    if(m_pyMatch.size() == 0)
    {
        std::cerr << "m_pyMatch.size() == 0, return false" << std::endl;
        return false;
    }

    if(m_depthCount == 0) return m_pyMatch[0].rotatedMatch(image, rrt, score);

    cv::Mat tempImage = image.clone();
    std::vector<cv::Mat> imagePyramid;
    imagePyramid.push_back(tempImage.clone());// 第 0 层就是原始图像。
    for(int i = 1; i <= m_depthCount; i++)  // 构造图像金字塔
    {
        cv::pyrDown(tempImage, tempImage);
        imagePyramid.push_back(tempImage.clone());
    }

    size_t depthCount = m_depthCount;
    cv::Rect rect;
    size_t index = 0;
    // 第一次是全局搜索。后面都是小区域内搜索
    bool ret = m_pyMatch[depthCount].rotatedMatch(imagePyramid[depthCount], rect, index, score);
//    std:: cout << "rotatedMatch " << endl;
    if(!ret) return false;
    do
    {
        depthCount --;
        index *= 2;
        rect.x *= 2; rect.y *= 2; rect.width *= 2; rect.height *= 2;

        m_pyMatch[depthCount].roi(imagePyramid[depthCount], index, 3, rect);
        // 到这里就获得了合适的 ROI 区域
        m_pyMatch[depthCount].regionMatch(imagePyramid[depthCount], rect, index, score);
    }while(depthCount != 0);

    double centerX = rect.x + rect.width / 2.0;
    double centerY = rect.y + rect.height / 2.0;

    rrt.center = cv::Point2f(centerX, centerY);
    rrt.size = rect.size();
    rrt.angle = -m_pyMatch[depthCount].angle(index);

    return true;
}

bool PyramidTemplateMatch::buildModels(const cv::Mat &model, const cv::Mat &mask)
{
    if(model.empty()) return false;
    if(model.size() != mask.size()) return false;

    m_model = model.clone();
    m_mask = mask.clone();
    //double angleMax = m_angleMin + m_angleCount * m_angleDet;
    //double angleDet = m_angleDet;
    m_depthCount = 0;
    cv::Mat model2 = model.clone(), mask2 = mask.clone();
    m_pyMatch.clear();
    int modelSize, angleCount = m_angleCount;
    do
    {
        Qly::RotateTemplateMatch2 match;
        m_pyMatch.push_back(match);
        //m_pyMatch[m_depthCount].setModelByDetAngle(model2, mask2, m_angleMin, angleMax, angleDet);
        m_pyMatch[m_depthCount].setModelByAngleCount(model2, m_angleMin, m_angleMax, angleCount);
        m_pyMatch[m_depthCount].setMethod(m_matchMethod);
        m_pyMatch[m_depthCount].setSizeMode(m_sizeMode);
        if(m_depthCount == 0) m_pyMatch[m_depthCount].setSubPixMode(m_enableSubPix);
        //std::cout << "m_depthCount = " << m_depthCount << ", Model size =" << Msg(model2) << std::endl;
        std::cout << "modelSize = " << Msg(model2) << ", m_depthCount = " << m_depthCount << ", angleCount = " << angleCount << std::endl;
        m_depthCount ++;
        cv::pyrDown(model2, model2);
        cv::pyrDown(mask2, mask2);
        angleCount = angleCount / 2;
        //angleDet *= 2;
        //angleCount = (angleCount + 1) / 2;
        //angleMax = m_angleMin + angleDet * angleCount;
        modelSize = std::min(model2.cols, model2.rows);

    }while(modelSize > m_minModelSize && m_depthCount <= m_maxDepth && angleCount >= m_minAngleCount);
    m_depthCount --; // 多加的那次要减下去
    std::cout << "PyramidTemplateMatch::setModelByAngleCount  final depthCount = " << m_depthCount << std::endl;
    return true;
}

bool RotateTemplateMatch2::regionMatch(const cv::Mat &image, cv::Rect &rect, size_t &index, double &score)
{
    if(index >= m_angleCount - 1) return false;
    if( rect.width == 0|| rect.height == 0 )
    {
        rect.x = 0;
        rect.y = 0;
        rect.width = image.cols;
        rect.height = image.rows;
    }

    size_t bestIndex;
    double bestScore;
    cv::Rect bestRect, roi = rect;

    switch (m_matchMethod) {
    case cv::TM_SQDIFF_NORMED:
    case cv::TM_SQDIFF:
        bestScore = std::numeric_limits<float>::max();
        break;
    default:
        bestScore = std::numeric_limits<float>::min();
        break;
    }
    bool ret1 =false , ret2 = false;
    double score1, score2;
    cv::Rect rect1, rect2;

#pragma omp parallel sections
{
#pragma omp section
    ret1 = singleMatch(image(roi), index, rect1, score1);
#pragma omp section
    ret2 = singleMatch(image(roi), index + 1, rect2, score2);
}
    if( ret1 && better(bestScore, score1, m_matchMethod))
    {
        bestScore = score1;
        bestRect = rect1;
        bestIndex = index;
    }
    if(ret2 && better(bestScore, score2, m_matchMethod))
    {
        bestScore = score2;
        bestRect = rect2;
        bestIndex = index + 1;
    }
    score = bestScore;
    rect = bestRect;
    rect.x += roi.x;
    rect.y += roi.y;
    index = bestIndex;
    return true;
}

bool RotateTemplateMatch2::regionMatch(const cv::Mat &image, const cv::Rect &rect, double angle, cv::RotatedRect &rrt, double &score)
{

    size_t index = static_cast<size_t>((angle - 0.5 * m_angleDet - m_angleMin) / m_angleDet);

//    std::cout << "IN RegionMatch, angleMin = " << m_angleMin << ", angleMax = " << m_angleMax << ", angleDet = " << m_angleDet << ", angleCount = " << m_angleCount  << std::endl;
//    std::cout << "---> angle = " << angle << ", angle - angleMin = " << angle - m_angleMin << ", index = " << index << std::endl;
    if(index >= m_angleCount - 1 || index < 0)
    {
        rrt.center = cv::Point2f(0, 0);
        rrt.size = cv::Size2f(0, 0);
        rrt.angle = 0;
        return false;
    }

//    std::cout << "m_rotatedModels[index] size= " << m_rotatedModels[index].size() << ", model size= " << m_model.size() << std::endl;
//    std::cout << "subpix = " << m_enableSubPix << endl;
    cv::Rect roi = rect;
    if( roi.width == 0|| roi.height == 0 )
    {
        roi.x = 0;
        roi.y = 0;
        roi.width = image.cols;
        roi.height = image.rows;
    }

    size_t bestIndex;
    double bestScore;
    cv::RotatedRect bestRRt;

    switch (m_matchMethod) {
    case cv::TM_SQDIFF_NORMED:
    case cv::TM_SQDIFF:
        bestScore = std::numeric_limits<float>::max();
        break;
    default:
        bestScore = std::numeric_limits<float>::min();
        break;
    }

    bool ret1 =false , ret2 = false;
    double score1, score2;
    cv::RotatedRect rrt1, rrt2;

#pragma omp parallel sections
{
#pragma omp section
    ret1 = singleMatch(image(roi), index, rrt1, score1);
#pragma omp section
    ret2 = singleMatch(image(roi), index + 1, rrt2, score2);
}
    if( ret1 && better(bestScore, score1, m_matchMethod))
    {
        bestScore = score1;
        bestRRt = rrt1;
        bestIndex = index;
    }
    if(ret2 && better(bestScore, score2, m_matchMethod))
    {
        bestScore = score2;
        bestRRt = rrt2;
        bestIndex = index + 1;
    }

    double bestAngle = m_angles[bestIndex]; // 这个是原始的计算值，可以用插值做改善

    if(m_enableSubPix)
    {
        if(bestIndex == 0 || bestIndex == m_angleCount - 2) return true; // 这时不用修正。直接返回

        if(bestIndex == index)
        {
            singleMatch(image(roi), index - 1, rrt, score);
        }
        else
        {
            singleMatch(image(roi), index + 2, rrt, score);
        }

        angleInterpolate(bestIndex, bestAngle, score);

        if( better(bestScore, score, m_matchMethod) )
        {
            bestScore = score;
        }
//        rrt.angle = -bestAngle;
    }

    score = bestScore;
    rrt = bestRRt;
    rrt.center.x += roi.x;
    rrt.center.y += roi.y;

    rrt.angle = bestAngle;

    return true;
}



bool RotateTemplateMatch2::roi(const cv::Mat &image, size_t angleIndex, int border, cv::Rect &rect)
{
    if(angleIndex >= m_angleCount - 1) return false;

    // 先计算 model 的最大尺寸。
    int width = m_rotatedModels[angleIndex].cols;
    int height = m_rotatedModels[angleIndex].rows;
    if(angleIndex != 0)
    {
        if(m_rotatedModels[angleIndex - 1].cols > width) width = m_rotatedModels[angleIndex - 1].cols;
        if(m_rotatedModels[angleIndex - 1].rows > height) height = m_rotatedModels[angleIndex - 1].rows;
    }
    if(angleIndex != m_angleCount - 1)
    {
        if(m_rotatedModels[angleIndex + 1].cols > width) width = m_rotatedModels[angleIndex + 1].cols;
        if(m_rotatedModels[angleIndex + 1].rows > height) height = m_rotatedModels[angleIndex + 1].rows;
    }

    width += border * 2;
    height += border * 2; // 搜索区域至少要有这么大

    if(width > rect.width)
    {
        rect.x -= (width - rect.width) / 2;
        rect.width = width;
        if(rect.x < 0) rect.x = 0;
        if( rect.x + width >= image.cols ) rect.x -= (rect.x + width - image.cols);
    }
    if(height > rect.height)
    {
        rect.y -= (height - rect.height) / 2;
        rect.height = height;
        if(rect.y < 0) rect.y = 0;
        if( rect.y + height >= image.rows ) rect.y -= (rect.y + height - image.rows);
    }
    // 至此，确定了搜索区域
    return true;
}

bool RotateTemplateMatch2::roi(const cv::Mat &image, const cv::RotatedRect &rrt, int border, cv::Rect &rect)
{
    int angleIndex;
    if(!index(rrt.angle, angleIndex))
    {
        rect.x = 0;
        rect.y = 0;
        rect.width = image.cols;
        rect.height = image.rows;
        return false;
    }

    // 先计算 model 的最大尺寸。
    int width = m_rotatedModels[angleIndex].cols;
    int height = m_rotatedModels[angleIndex].rows;
    if(angleIndex != 0)
    {
        if(m_rotatedModels[angleIndex - 1].cols > width) width = m_rotatedModels[angleIndex - 1].cols;
        if(m_rotatedModels[angleIndex - 1].rows > height) height = m_rotatedModels[angleIndex - 1].rows;
    }
    if(angleIndex != m_angleCount - 1)
    {
        if(m_rotatedModels[angleIndex + 1].cols > width) width = m_rotatedModels[angleIndex + 1].cols;
        if(m_rotatedModels[angleIndex + 1].rows > height) height = m_rotatedModels[angleIndex + 1].rows;
    }

    width += border * 2;
    height += border * 2; // 搜索区域至少要有这么大

    rect.x = rrt.center.x - width / 2;
    rect.y = rrt.center.y - height / 2;
    rect.width = width;
    rect.height = height;

    if(rect.x < 0) rect.x = 0;
    if( rect.x + width >= image.cols ) rect.x -= (rect.x + width - image.cols);
    if(rect.y < 0) rect.y = 0;
    if( rect.y + height >= image.rows ) rect.y -= (rect.y + height - image.rows);

    return true;

}

bool RotateTemplateMatch2::angle(size_t index, double &angle)
{
    if(index >= m_angleCount) return false;
    angle = m_angles[index];
    return true;
}

double RotateTemplateMatch2::angle(size_t index)
{
    if(index >= m_angleCount) return -361.0;

//    if(m_enableSubPix)
//    {
//        double bestAngle, score;
//        angleInterpolate(index, bestAngle, score);
//        return bestAngle;
//    }
    return m_angles[index];
}

bool RotateTemplateMatch2::index(double angle, int &idx)
{
    idx = static_cast<int>((angle - m_angleMin) / m_angleDet - 0.5);
    if(idx < 0 || idx > m_angleCount - 1) return false;
    return true;
}


}
